package com.xixi.core.rpc.registry;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.collection.ConcurrentHashSet;
import cn.hutool.cron.CronUtil;
import cn.hutool.cron.task.Task;
import cn.hutool.json.JSONUtil;
import com.xixi.core.rpc.config.RegistryConfig;
import com.xixi.core.rpc.model.ServiceMetaInfo;
import io.etcd.jetcd.*;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.watch.WatchEvent;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

/**
 * @author: zxh
 * @date: 2024/8/22 23:23
 * @description: Etcd 注册中心实现
 */
public class EtcdRegistry implements Registry {

	private Client client;

	private KV kvClient;

	// 根节点
	private static final String ETCD_ROOT_PATH = "/rpc/";

	// 本机注册的节点key集合，用于维护续期
	private final Set<String> localRegisterNodeKeySet = new HashSet<>();

	// 注册中心服务缓存
	private final RegistryServiceCache registryServiceCache = new RegistryServiceCache();

	// 正在监听的key集合
	private final Set<String> watchingKeySet = new ConcurrentHashSet<>();

	@Override
	public void init(RegistryConfig registryConfig) {
		client = Client.builder()
				.endpoints(registryConfig.getAddress())
				.connectTimeout(Duration.ofMillis(registryConfig.getTimeout()))
				.build();
		kvClient = client.getKVClient();
	}

	@Override
	public void registry(ServiceMetaInfo serviceMetaInfo) throws ExecutionException, InterruptedException {
		// 创建 Lease 和 KV 客户端
		Lease leaseClient = client.getLeaseClient();

		// 创建一个30秒的租约
		long leaseId = leaseClient.grant(30).get().getID();

		// 设置要存储的键值对
		String registerKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
		System.out.println("注册：" + registerKey);
		ByteSequence key = ByteSequence.from(registerKey, StandardCharsets.UTF_8);
		ByteSequence value = ByteSequence.from(JSONUtil.toJsonStr(serviceMetaInfo), StandardCharsets.UTF_8);

		// 将键值对与租约关联起来，并设置过期时间
		PutOption putOption = PutOption.builder()
				.withLeaseId(leaseId)
				.build();
		kvClient.put(key, value, putOption).get();

		// 添加节点信息到本地缓存
		localRegisterNodeKeySet.add(registerKey);
	}

	@Override
	public void unRegistry(ServiceMetaInfo serviceMetaInfo) throws ExecutionException, InterruptedException {
		String registerKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
		System.out.println("注销：" + registerKey);
		kvClient.delete(ByteSequence.from(registerKey, StandardCharsets.UTF_8)).get();

		// 注销时移出缓存的key
		localRegisterNodeKeySet.remove(registerKey);
	}

	@Override
	public List<ServiceMetaInfo> serviceDiscovery(String serviceKey) {
		// 优先从缓存获取服务
		List<ServiceMetaInfo> cacheServiceInfo = registryServiceCache.readCache();
		if (cacheServiceInfo != null) {
			return cacheServiceInfo;
		}
		// 前缀搜索，结尾一定要加 '/'
		String searchPrefix = ETCD_ROOT_PATH + serviceKey + "/";

		try {
			// 前缀查询
			GetOption getOption = GetOption.builder()
					.isPrefix(true)
					.build();
			List<KeyValue> keyValues = kvClient.get(ByteSequence.from(searchPrefix, StandardCharsets.UTF_8),
							getOption)
					.get()
					.getKvs();
			// 解析服务信息
			List<ServiceMetaInfo> serviceMetaInfoList = keyValues.stream().map(keyValue -> {
				String key = keyValue.getKey().toString(StandardCharsets.UTF_8);
				// 监听 key
				watch(key);
				String value = keyValue.getValue().toString(StandardCharsets.UTF_8);
				return JSONUtil.toBean(value, ServiceMetaInfo.class);
			}).collect(Collectors.toList());

			// 写入服务缓存
			registryServiceCache.writeCache(serviceMetaInfoList);
			return serviceMetaInfoList;
		} catch (Exception e) {
			throw new RuntimeException("获取服务列表失败：", e);
		}
	}


	@Override
	public void destroy() {
		System.out.println("当前节点下线");
		// 下线节点
		// 遍历本节点的所有key
		for (String key : localRegisterNodeKeySet) {
			try {
				kvClient.delete(ByteSequence.from(key, StandardCharsets.UTF_8)).get();
			} catch (Exception e) {
				throw new RuntimeException(key + "节点下线失败");
			}
		}
		// 释放资源
		if (kvClient != null) {
			kvClient.close();
		}
		if (client != null) {
			client.close();
		}
	}

	@Override
	public void heartBeat() {
		System.out.println("开启心跳检测");
		// 10 秒续签一次
		CronUtil.schedule("*/10 * * * * *", (Task) () -> {
			// 遍历本节点的所有key
			for (String key : localRegisterNodeKeySet) {
				try {
					List<KeyValue> keyValues = kvClient.get(ByteSequence.from(key, StandardCharsets.UTF_8))
							.get()
							.getKvs();
					// 该节点已过期
					if (CollUtil.isEmpty(keyValues)) {
						continue;
					}
					// 节点未过期
					KeyValue keyValue = keyValues.get(0);
					String value = keyValue.getValue().toString(StandardCharsets.UTF_8);
					ServiceMetaInfo serviceMetaInfo = JSONUtil.toBean(value, ServiceMetaInfo.class);
					registry(serviceMetaInfo);
				} catch (Exception e) {
					throw new RuntimeException(key + "续签失败", e);
				}
			}
		});

		// 支持秒级别定时任务
		CronUtil.setMatchSecond(true);
		CronUtil.start();
	}

	@Override
	public void watch(String serviceNodeKey) {
		Watch watchClient = client.getWatchClient();
		// 之前未被监听，开启监听
		boolean newWatch = watchingKeySet.add(serviceNodeKey);
		if (newWatch) {
			watchClient.watch(ByteSequence.from(serviceNodeKey, StandardCharsets.UTF_8), watchResponse -> {
				for (WatchEvent event : watchResponse.getEvents()) {
					switch (event.getEventType()) {
						// key删除时触发
						case DELETE:
							// 清理注册服务缓存
							registryServiceCache.clearCache();
							break;
						case PUT:
						default:
							break;
					}
				}
			});
		}
	}
}
